Introduction

Our paper is "Unifying Heterogeneous Electronic Health Records Systems via Text-Based Code Embedding" by Hur et al. (2022). All paper references are provided at the bottom of this notebook.

While EHR (Electronic Health Records) provide an attractive data source to do research on exposures and disease, there are many heterogenous medical code formats used by different healthcare providers. Thus, clinical studies using EHRs are difficult to scale up due to data incompatibilities. The techni- cal hurdle is that previous systems learned hidden representations (latent embeddings) for each code system which projected the same medical concept with different encodings into different, incompatible semantic spaces. The paper by Hur et al (2022) learns from the unstructured textual descriptions instead of the medcodes. To this end the authors propose a novel method to unify heterogeneous EHR systems via a novel method to learn text embeddings of medical codes from their textual descriptions, which they call description embedding (DescEmb). themselves, a strategy which the authors termed description embedding (DescEmb). Using DescEmb the authors were able to pool differently structured EHRs from two datasets, namely MIMIC-III and eICU, into one larger pooled dataset and achieved higher accuracy. The authors then tested the performance of the DescEmb on five tasks: mortality prediction, length of stay prediction, readmission prediction and diagnosis. Hur et al. were to our knowledge the first to attempt a unification of different EHR systems via text-based code embeddings. The authors achieved state-of-the-art performance on all five tasks, which is a significant contribution to the field of EHR research.

Scope of Reproducibility:

We plan to focus on hypotheses 1 - 3 mainly. Hypothesis 4 is addressed by comparing the results of the proposed method with the state-of-the-art methods (see Results section). Hypothesis 5 is our own hypothesis which we will test by using a BERT model pre-trained on a domain-specific dataset. We view this as an extension of the original paper which is optional, thus it will not be included in the draft.

  1. Hypothesis 1: DescEmb outperforms CodeEmb
  2. Hypothesis 2: Two EHR datasets with different structure can be used interchangeably with the proposed method.
  3. Hypothesis 3: Differently structured EHR datasets can be pooled together to improve the performance of the model.
  4. Hypothesis 4: DescEmb has a competitive or superior performance compared to the state-of-the-art methods.
  5. Hypothesis 5 (ours): Using a BERT model which was pre-trained in a strictly domain-specific manner can improve the performance of the model further (not shown by the authors).

GitHub repository:

If you entered through a PDF, you can find our GitHub repository here.

Video

We uploaded our video walkthrough to YouTube.

Methodology

Environment Setup

We provide a script to create a virtual environment and install necessary packages.

./setup.sh

Imports

Prediction tasks studied in the paper

The authors assess model performance using five different prediction tasks listed in the following table.

Task Number Task identifier Task Description Classification type
Task 1 Dx Predicting the diagnosis of a patient given the clinical notes. Multi-label
Task 2 Mort Predicting the mortality of a patient given the clinical notes and the treatment codes. Binary
Task 3 LOS>3 Predicting whether the length of stay of a patient is greater than 3 days given the clinical notes and the treatment codes. Binary
Task 4 LOS>7 Predicting whether the length of stay of a patient is greater than 7 days given the clinical notes and the treatment codes. Binary
Task 5 ReAdm Predicting whether a patient will be readmitted within 30 days given the clinical notes and the treatment codes. Binary

For task 1 the highest level representation of ICD-9-CM is used, in total 18 representations. MIMIC-III already uses ICD-9, while eICU uses ICD-10. The authors map ICD-10 to ICD-9 using the Clinical Classificatons Software (CCS).

Abbreviations used in the paper

Abbreviation Description
$p^i$ The $i$-th patient
$c_i$ The $i$-th medical event (e.g. diagnoses, prescriptions)
$\mathcal{C}$ The set of all medical events
$t_i$ The $i$-th time stamp
$w_{i,j}$ The $j$-th word in the $i$-th medical event
$\mathcal{W}$ The set of all words, i.e. vocabulary of description
$d^i$ The $i$-th text description of diagnosis of the $i$-th patient
$E_{\Psi}$ The embedding layer of the medical events used in CodeEmb
$\boldsymbol{c}_i$ Embedded representation of the $i$-th medical event
$B_{\Phi}$ The embedding layer of the text descriptions used in DescEmb
$\boldsymbol{z}_i$ Embedded representation of the $i$-th text description

Data

We downloaded two datasets, MIMIC-III v1.4 from PhysioNet and eICU v2.0 from MIT Lab and downloaded it via PhysioNet using wget. The datasets are stored in the following paths in our Google Drive:

Data description

eICU consists of 200,859 patients, 1,139,695 admissions, and 2,839,547 patient unit stays. It is organized in three levels:

A comprehensive description of the dataset can be found in the paper by Pollard et al.

MIMIC-III consists of 46,520 patients, 58,976 admissions, and 7,875,529 clinical notes. An overview of the dataset:

A comprehensive description of the dataset can be found in the paper by Johnson et al.

Preprocessing

We preprocessed the data using the code provided by the authors in their GitHub repository. We ran

INPUT_PATH=/home/data
OUTPUT_PATH=/home/data/output
DX_PATH=$INPUT_PATH/ccs_multi_dx_tool_2015.csv

python ../preprocess/preprocess_main.py \
    --src_data mimiciii \
    --dataset_path $INPUT_PATH/mimic \
    --ccs_dx_tool_path $DX_PATH \
    --dest_path $OUTPUT_PATH ;

In total this was run six times: For three datasets, i.e. mimiciii, eicu or pooled and $data_type was either predict or pretrain - see our modified preprocess_run.sh for all six runs. This script needs to be run from within the author's DescEmb repository, because of all the dependencies which main.py imports. Preprocessing was done on a cloud node with 128 GB of RAM, 32 cores, RTX A6000 GPU and 180 GB of storage.

Both, data cleaning and cohort construction are done by running the same script. Below we describe what is happening in those steps.

The datasets are constructed such that three sources of information are available:

Each of these sources is available in two different datasets: MIMIC-III and eICU. The authors use the following files from the MIMIC-III and eICU datasets to construct the datasets for the tasks:

Item Source Filnename
Laboratory results MIMIC-III labevents.csv
Medication MIMIC-III prescriptions.csv
Infusion MIMIC-III inputevents_cv.csv, inputevents_mv.csv -> Merged together
Laboratory results eICU lab.csv
Medication eICU medication.csv
Infusion eICU infusionDrug.csv

Data Cleaning

On merging inputevents_cv.csv and inputevents_mv.csv, the authors found that the two datasets have different structures. They used the following steps to clean the data:

Cohort Construction

The authors used the following steps to construct the cohorts:

  1. Cohort: MICU patients whose first care unit was also the last care unit and of type ICU.
  2. Cohort: Patients with multiple ICU stays above age 18
  3. Cohort: Patients with multiple ICU stays who remained in ICU for > 12 hours

Datasets and Dataloaders

In the proposal we intended to leverage pyhealth to load the data. However, we found that the authors' code for preprocessing may not be compatible with pyhealth. We will use the authors' code to load the data.

Four dataset classes available in datasets/dataset.py, all of which inherit from BaseDataset, which in turn inherits from torch.utils.data.Dataset. The four classes are:

The base class BaseDataset provides the following functionality:

Class Description Trainer class embed_model
CodeDataset Dataset for the code embeddings Trainer codeemb
TokenizedDataset Dataset for the tokenized descriptions Trainer descemb
MLMTokenizedDataset Dataset for the masked language model tokenized descriptions Trainer mlm
Word2VecDataset Dataset for the word2vec tokenized descriptions Word2VecTrainer w2v

CodeDataset returns stacked tensors comprising input ID's, sequence length, value and label.

TokenizedDataset stacks input IDs, token type IDs, attention masks, sequential lengths, values, and labels into tensors and organizes them into a dictionary, ready for model input.

MLMTokenizedDataset returns input IDs, token type IDs, attention masks, MLM labels for masked language model training.

Word2VecDataset returns an index dictionary as well as positive and negative word pairs to train the word2vec model usin a skip-gram approach.

The datasets are utilized depending on which embedding model is being trained. This decision is made by the user as a command-line argument to main.py. There are two trainer classes, one for word2vec and one for all other cases. The table above shows which dataset is used for which embedding model.

Exploratory Data Analysis

  1. Load the data
  2. Check the head of the data
  3. Check the missing values in the data
  4. Check the value counts in the data
  5. Check demographics of the patients
Load the data
Check the head of the data
Check the missing values in the data
Check value counts in the data
Check demographics of the patients

Model

The model includes the model definitation which usually is a class, model training, and other necessary parts.

CodeEmb vs. DescEmb strategies

The authors propose two strategies to embed the data: CodeEmb and DescEmb. The two strategies are described as follows:

$$ \begin{align*} \textit{Code-based Embedding:} \\ &\boldsymbol{c}_i = E_{\Psi}(\boldsymbol{c}_i) \\ &\hat{y} = P_{\Phi}(\boldsymbol{c}_1, \boldsymbol{c}_2, ..., \boldsymbol{c}_T) \\ \textit{Description-based Embedding:} \\ &d_i = (w_{i,1}, w_{i,2}, ... , w_{i,n}) \\ &\boldsymbol{z}_i = B_{\Psi}(d_i) \\ &\hat{y} = P_{\Phi}(\boldsymbol{z}_1, \boldsymbol{z}_2, ..., \boldsymbol{z}_T ) \end{align*} $$

Effectively, the CodeEmb strategy uses the medical events directly, while the DescEmb strategy uses the text descriptions of the medical events. The authors use a Bi-LSTM for the embedding layer $E_{\Psi}$ in CodeEmb. For $B_{\Psi}$ they ended up using a BERT (Bi-directional Encoder Representations from Transformers). While the authors tested several model depths, we will only use the smallest 2-MSA-layer model BERT-tiny. For the most part $\boldsymbol{z_i}$ corresponds to the BERT output vector for the [CLS] token. The prediction layer $P_{\Phi}$ employs an RNN.

Value Embedding

The descriptions of medical events often contain numbers (values), which are important for the prediction tasks. The authors propose four different methods to embed these values into the model. The methods is described as follows:

The following figure summarizes the value embedding techniques.

Model Training

There are three optimization schemes used in the paper:

$$ \begin{align} \textit{CodeEmb}: argmin_{\Theta, \Psi} \mathcal{L}(\hat{y}, y) \\ \textit{DescEmb, all parameters}: argmin_{\Theta, \Phi} \mathcal{L}(\hat{y}, y) \\ \textit{DescEmb, class fine-tuning}: argmin_{\Psi, \boldsymbol{z}_{CLS}} \mathcal{L}(\hat{y}, y) \\ \end{align} $$

Eqn. 1 is the optimization scheme for CodeEmb, where $\Theta$ are the parameters of the prediction layer and $E_{\Psi}$, and $\Psi$ are the parameters of the embedding layer. Eqn. 2 is the optimization scheme for DescEmb where all parameters are fine-tuned. Eqn. 3 is the optimization scheme for DescEmb where only the class token $\boldsymbol{z}_{CLS}$ is fine-tuned, but the parameters of the text embedding layer $E_{\Psi}$ are fixed.

We used the author's example scripts for running model training and provide our scripts in the project_code folder.

Script Description Essential Arguments
00_pretrain_codeemb.sh Pretrain the CodeEmb model (i.e. Embedding layer followed by RNN) --model ehr_model --embed_model codeemb --value_mode NV --task w2v
01_pretrain_descemb_rnn.sh Pretrain the DescEmb encoder with RNN architecture and MLM target --model descemb_rnn --value_mode NV --task mlm
02_pretrain_descemb_bert.sh Pretrain the DescEmb encoder with BERT architecture and MLM target --model descemb_bert --value_mode NV --task mlm
00_single_domain_learning_descemb.sh Train the DescEmb model with fine-tuning on prediction task Iterate through --embed_models=('descemb_rnn' 'descemb_bert'), --tasks=('readmission' 'mortality' 'los_3day' 'los_7day' 'diagnosis') and --value_modes=('NV' 'VA' 'DSVA' 'DSVA_DPE' 'VC')
01_single_domain_learning_codeemb.sh Train the CodeEmb model with fine-tuning on prediction task Iterate through --embed_models=('codeemb'), --tasks=('readmission' 'mortality' 'los_3day' 'los_7day' 'diagnosis') and --value_modes=('NV' 'VA' 'DSVA' 'VC')

As we progressed with the project the number of scripts we produced grew, but they all followed the same overall logic. You can find a comprehensive list in the project_code folder.

For illustration we reproduce one of them here:

#!/bin/bash
# Run training new model using FT with descemb
# run this from project root folder using
# ./project_code/02_train_new_models/03_train_ft_descemb_bert.sh $(pwd)/project_code

if [ -n "$1" ]; then
    root="${1%/}/DescEmb/"
    MODEL_PATH=("${1%/}/DescEmb/outputs/2024-04-19/06-18-04/checkpoints/checkpoint_mimiciii_descemb_bert_mlm_last.pt")
else
    root='../DescEmb/'
    MODEL_PATH=('../../../outputs/2024-04-19/06-18-04/checkpoints/checkpoint_mimiciii_descemb_bert_mlm_last.pt')
fi 

INPUT_PATH=/home/data/output
SRC_DATA=('mimiciii' 'eicu')

embed_models=('descemb_bert')
tasks=('diagnosis' 'los_3day' 'los_7day' 'readmission' 'mortality')
value_modes=('VA' 'DSVA' 'DSVA_DPE' 'VC')


for data in "${SRC_DATA[@]}"; do
    for index in "${!embed_models[@]}"; do
        emb_model=${embed_models[$index]}
        model_path=${MODEL_PATH[$index]}
        echo "Processing model: $emb_model with data $data"
        for task in "${tasks[@]}"; do
            for value_mode in "${value_modes[@]}"; do
                echo "Current directory: $(pwd)" && \
                CUDA_VISIBLE_DEVICES=0 python ${root}main.py \
                    --distributed_world_size 1 \
                    --input_path "$INPUT_PATH" \
                    --init_bert_params \
                    --model ehr_model \
                    --embed_model "$emb_model" \
                    --pred_model rnn \
                    --src_data $data \
                    --ratio 100 \
                    --patience 90 \
                    --value_mode "$value_mode" \
                    --save_prefix "checkpoint_${data}_${emb_model}_${task}" \
                    --task "$task" ;
            done
        done
    done
done

Because we modified the author's code in several places to remove issues we had with preprocessing, GPU utilization and others we kept a copy in a subfolder called DescEmb. The training loop is part of the trainer classes in the trainer.py file. The training loop is a standard PyTorch training loop. We did modify these scripts to run on a single GPU, as we only have access to a single GPU. The scripts provided by the authors are for multi-GPU training.

## Before the training loop get the available GPU devices (exerpt)
self.n_devices = torch.cuda.device_count() if torch.cuda.is_available() else 0
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if self.n_devices > 1:
    self.model = nn.DataParallel(model, device_ids=args.device_ids).to(self.device)
else:
    self.model = model.to(self.device)
## continue trainer init ...

## Inside the training loop (exerpt)
for sample in tqdm.tqdm(self.data_loaders['train']):
    self.optimizer.zero_grad(set_to_none=True)
    net_output = self.model(**sample["net_input"])
    if self.n_devices > 1:
        logits = self.model.module.get_logits(net_output)
        target = self.model.module.get_targets(sample).to(logits.device)
    else:
        logits = self.model.get_logits(net_output)
        target = self.model.get_targets(sample).to(logits.device)
## continue loop ...

The number of hyperparameters needed to uniquely determine each experiment is at least 7, plus some additional parameters regarding data and model paths. For this reason we produced the following table to generate an overview of the hyperparameters used in the paper. We also use this to keep track of which experiments need to be run in order to reproduce the results from the paper. assets/Overview_training_conditions.png

The pink cells are the key (combinations) which determine what kind of experiment it is. E.g. for CLS-FT you need to use an ehr_model and init_bert_params_with_freeze.

Results and Analysis

Preprocessing

We had many issues during the preprocessing step, because not all options are documented. So, at first we learned them from the command-line options in the source code. Then the authors released an updated version which fixed some of the issues we had (which we are very appreciative of). There was still a lot of trial and error involved. Still on the the last day before the draft deadline the authors were still making changes.

Some issues were related to the folder structure the scripts are creating to deviate from what some parts of the processing are expecting. We filed an issue on that. As a workaround we used some manual copying and soft linking of files to the best of our understanding. However, it is likely that we made some mistakes in this process. Still for the draft submission we utilized the preprocessed data thus obtained. After that things had stabilized. Unfortunately, one issue remained invisible to us until very late in the project: pooled training was missing some crucial files with DSVA_DPE value embedding. More details further below.

Pretraining DescEmb with RNN or BERT encoders

In order to get DescEmb to run we need to pretrain the text encoders. We ran the scripts provided by the authors for pretraining the DescEmb model with RNN and BERT encoders. We ran the scripts for pretraining the DescEmb model with RNN and BERT encoders. The scripts are provided in the project_code folder (see table above). A serious issue we had has been that all pretrainings with --task mlm got stalled. Debugging showed, that there must be an issue with the MLMTokenizedDataset class. Apparently, there is no noteable progress when loading the first batch of data, even when letting it run for hours (on the same machine we used for pre-processing). At one point we hypothesized that maybe the dataset is too large and created a minimal version MIMIC-III with only 1% of the data randomly sampled. We used 03_Creating_minimal_dataset.ipynb to that end. This did not solve our problem either. Eventually, we found the root cause (see Discussion / Issues summary).

Experiment Tracking

There is a substantial number of experiments to run if we wish to reproduce the results of the paper. Each of Tables 1 and 2 in the paper reports the AUPRC on five prediction tasks with 3-4 different value embedding methods and different models. This results in 170 experiments to run each. Table 3 reports the AUPRC on single domain learning, transfer learning and pooled learning, which amounts to 120 experiments. We need to keep track not only of the AUPRC metrics, but also the hyperparameters used in each experiment. Likewise, if a run doesn't finish, we need to keep track of the reason why. While popular tools such as Weights & Biases or mlflow exist for this purpose, they, too require some overhead to set up. We will use a simple logging system for now: The Trainer class in the trainer.py file has a log method which logs the loss and AUPRC progression, as well as hyperparameters to a file called train.log. We will use this to log the hyperparameters and the results of each experiment.

In order to generate and overview as well as results table we produced some helper functions to parse the results from these log files. We provide these functions in the project_code/evaluate folder in logs.py. Some more helper functions are in the works to generate training history plots and other visualizations, much like tensorboard would do, but simpler and more tailored to our needs.

A global table with all experiments is stored as an Excel file in the project_code/outputs/experiments.xlsx file. All training logs and hyperparameters for all runs are documented in project_code/DescEmb/outputs. We did not upload all trained models to GitHub, but zipped them and put them on a shared folder in Google Drive.

We create the experiments.xlsx and metrics.xlsx by running python parse_experiments.py in the root folder of this project. The script also applies the overview table of all experiment types shown in the Methods section to introduce a tag column. Typical output of the script looks like:

Found 549 experiments and 105321 metrics.
CodeEmb_RD                   94
DescEmb-BERT_FT              90
DescEmb-RNN_Scr              78
DescEmb-RNN_Scr-MLM          75
DescEmb-BERT_FT-MLM          69
DescEmb-BERT_CLS-FT          46
DescEmb-BERT_Transfer        32
DescEmb-BERT_Scr             24
DescEmb-RNN_Transfer         24
DescEmb-RNN_Pretrain-MLM      3
DescEmb-BERT_Pooled           3
CodeEmb-Pooled                3
unknown                       2
DescEmb-BERT_Pretrain-MLM     2
CodeEmb_Pretrain-W2V          2
CodeEmb_W2V                   2

Model Performance

The authors reported mainly AUPRC, i.e. area under the precision-recall curve, as the evaluation metric. There are three main tables in the paper which we will try to reproduce.

Hypothesis 1

We need to reproduce (parts of) Tables 1 and 2 in the paper. We will use the AUPRC metric to evaluate the performance of the models. We will compare the performance of the CodeEmb and DescEmb strategies on the five tasks with and without pretraining as well as different value embedding strategies.

Tracking of Experiments

Each run produces a log file which contains the hyperparameters and the metrics (loss, auroc, auprc) of the run. We parse these log files to generate an overview of the experiments and the results. We use this to generate the tables for the paper. Additionally, we use the log files to create diagnostic plots to understand the training process better. In the run below we noticed overfitting after about 300 epochs. Overfitting shows in two ways: The validation loss starts to increase and the validation AUPRC starts to decrease. During training we used early stopping to prevent overfitting. We kept the model with the best (lowest) validation loss.

MIMIC-III

eICU

Hypothesis 1 asks if DescEmb has higher predictive power than direct EHR code embedding using CodeEmb. We agree this is largely the case. We take note of the following:

For the following experiments the authors used CodeEmb-RD + VC and (DescEmb-BERT-FT, DescEmb-RNN-Src, DescEmb-RNN-Src-MLM) + DSVA_DPE. Since these are mostly not available, but performance of DSVA and DSVA_DPE were very similar we continued using DSVA value encoding.

Hypothesis 2 and 3: Single Domain Learning vs. Transfer learning

Embedding Visualization

Embedding Takeaways

Above are UMAP plots of the various model embeddings. UMAP is a dimensionality reduction technique that aims to preserve the local and global structure of high-dimensional data in a low-dimensional space. When we reduce the high dimensionional embedding space, we are able to see the underlying structure and relationship the data points have with each other. Since most of the tasks are binary classification problems, we should see 2 clusters in many tasks which show that the training model was able to learn embeddings that created a separable space for the two classes. While the labels are not present and are non trivial to gather, we are still able to see the seprability in datapoints in the embedding space. We use that to get an idea about the class sepration in embedding space. A clear separation is indicative of a good classifier.

The los 7 day and los 3 day tasks consistently separate the best, and readmission task performs the worst with the embedding space looking like a gaussian blob in both the MIMIC and eICU datasets.

It is also interesting to note that the eICU separates better than the MIMIC dataset just from looking at the embeddings, which is consistent with the authors findings and our findings from the AUPRC tables above.

Discussion

Hypothesis 1

Overall we confirm the main observation of Tables 1 and 2 in the paper in that DescEmb achieves higher AUPRC than CodeEmb. We did not quite reach the same AUPRC values as reported in the paper. The reason for that is not known to us. For training we used the default parameters suggested by the author's main.py. We hypothesized, however, that some training runs might have been run for too few epochs. So, we used different patience settings.

While we can't show the training history for all conditions we picked DescEmb-BERT-FT DSVA_DPE trained on both datasets with patience settings 5, 45, 90 and for two tasks, namely readmission and diagnosis prediction. We notice:

Another indication of overfitting is that the test results tend to perform worse than validation, indicative of poor generalization of the model. Only our train results come close to what the authors reported.

Hypothesis 4 - DescEmb outperforms CodeEmb in predicting patient outcomes

The authors report that DescEmb outperforms CodeEmb in most tasks. The basis for which they make this claim is the AUPRC metric on five tasks: Dx, Mort, LOS>3, LOS>7, ReAdm. Wornow et al. (2023) criticize the use of AUPRC as the sole metric for evaluation in combination with a narrow set of prediciton tasks. They argue that the tasks used were not always overlapping with other publications thus making it hard to compare the results. Wornow et al. (2023) provide a taxonomy of medical foundation models. According to this taxonomy the tasks used in this paper would fall under the category of "prediction tasks" by an FEMR.

While in NLP tasks standardized benchmarks exist, such as BLUE and ROUGE the same is not always true in the medical domain. However, there have been some efforts to standardize the evaluation of medical NLP tasks, such as the BLURB benchmark put forward by Gu et al. (2022). To the best of our knowledge the same is unfortunately not true for the tasks used in this paper.

Hypothesis 5 - Using a BERT model pre-trained on a domain-specific dataset (optional)

To assess if such a domain-specific pre-training improve the can an improve performance we had planned to replace the AutoTokenizer in the BaseDataset by one derived from PubMedBERT as reported by Gu et al (2022). Then we could run preprocessing, pre-training and pooled training with this tokenization. Unfortunately, we did not have enough time to do all these steps, due to other, more basic issues as discussed.

Issues and Challenges

To start with we really appreciate the authors made comprehensive code available. This is a great help for reproducibility. However, the code has had some issues when we started out. We also had to understand all the available command-line options by reading both documentation and source code as well as a lot of experimentation.

Summary of issues we experienced

Issue Impact Solution
Bugs in preprocessing when we started Could not train any DescEmb model Mostly solved by author's fixes made in April 2024
Data with DSVA_DPE value mode not generated in preprocessing Could not do pooled training Modification to preprocess_main.py (too) late in the project
Pretraining with MLM target took excessively long MLM pretraining not possible We narrowed it down to MLMTokenizedDataset and made modifications to it (see below)
Training CodeEmb with VC value mode did not run We could not include CodeEmb in Table 3 which used VC None so far, we used DSVA instead
Training CodeEmb-W2V did not run Less fair comparison between CodeEmb and DescEmb None so far
Transfer learning DescEmb-RNN with DSVA_DPE did not run Can not include in Table 3 None so far, we used DSVA instead
Pooled training didn't run Can't speak to how pooling data impacts model performance None so far. We solved part the issues with preprocessing, but ultimately, training still failed

Changes made to MLMTokenizedDataset

The challenges we had were such that cuda was not properly initialized. This may have been due to the fact that the authors trained on a multi-GPU machine and used nn.DataParallels which handles GPU initialization diferently. So, we had to assign tensors to GPUs in several places within the dataset class, for example:

input = {"input_ids": torch.stack([s["input_ids"] for s in samples]).to(self.device)}
        input["seq_len"] = torch.stack([s["seq_len"] for s in samples]).to(self.device)
        input["value"] = torch.stack([s["value"] for s in samples]).to(self.device)
        out = {"label": torch.stack([s["label"] for s in samples]).to(self.device)}

Here self.device was

self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

Given the age of the publication the authors used nn.DataParallels for multi-GPU training, which is now deprecated in PyTorch. We did NOT try to update the code to use torch.nn.parallel.DistributedDataParallel instead, because we wished to avoid multi-GPU training due to cost considerations.

Challenges

Generally, the sheer number of experiments is very large (175 per dataset for tables 1 and 2, 80 additional ones for table 3) and the amount of RAM and VRAM necessary partially went beyond of what we could run on our local machines (mostly for preprocessing, pretraining and some variants of DescEmb-BERT with FT and FT-MLM). The memory limitation was overcome by renting a cloud node with 1 RTX6000 Ada GPU that had 48 GB of VRAM, an AMD EPYC 7713 64-Core Processor and 1 TB of RAM (the latter was not fully used). Over the course of the experiments we used up about 180 GB of storage, which was just about enough (when deleting some experiments we had downloaded from the server).

As for the number of experiments we resorted to bash scripting of the author's main.py. We used nested loops whereever possible to run a whole batch of experiments, e.g. DescEmb-BERT-FT with both datasets, all five tasks and all four value modes. The scripts are part of our repository. For the mentioned example see here. To keep track of the experiments we wrote a parse_experiments.py which recursively parsed all runs in output folders and extracted the most important hyperparameters. It also added a tag column which helped us to quickly filter what kind of experiment each row corresponds to. This was particularly useful to manage contributions from a distributed team working asynchronously.

We did not train BERT-tiny from scratch, because it is very computationally expensive and the authors found that it did not help very much over fine-tuning.

References

  1. Kyunghoon Hur, Jiyoung Lee, Jungwoo Oh, Wesley Price, Young-Hak Kim, and Edward Choi. Uni- fying Heterogeneous Electronic Health Records Systems via Text-Based Code Embedding, March 2022, arXiv:2108.03625, doi: 10.48550/arXiv.2108.03625.
  2. Pollard, Tom, et al. "eICU Collaborative Research Database" (version 2.0). PhysioNet (2019), https://doi.org/10.13026/C2WM1R.
  3. Johnson, Alistair, et al. "MIMIC-III Clinical Database" (version 1.4). PhysioNet (2016), https://doi.org/10.13026/C2XW26.
  4. Wornow, Michael, Yizhe Xu, Rahul Thapa, Birju Patel, Ethan Steinberg, Scott Fleming, Michael A. Pfeffer, Jason Fries, and Nigam H. Shah. "The Shaky Foundations of Clinical Foundation Models: A Survey of Large Language Models and Foundation Models for EMRs." March 2023, arXiv:2303.12961, doi: 10.48550/arXiv.2303.12961.
  5. Gu, Yu, Robert Tinn, Hao Cheng, Michael Lucas, Naoto Usuyama, Xiaodong Liu, Tristan Naumann, Jianfeng Gao, and Hoifung Poon. "Domain-Specific Language Model Pretraining for Biomedical Natural Language Processing." January 2022, arXiv:2007.15779, doi: 10.1145/3458754.